﻿#pragma once
#include<iostream>
using namespace std;

enum Colo
{
	RED,
	BLACK
};

template<class T>
struct RBtreeNode
{
	T _data;               //存放数据
	RBtreeNode<T>* _left;  //左节点
	RBtreeNode<T>* _right; //右节点
	RBtreeNode<T>* _parent;//父亲节点
	Colo _col;//颜色

	RBtreeNode(const T& data)
		:_data(data)
		,_left(nullptr)
		,_right(nullptr)
		,_parent(nullptr)
	{

	}
};


template<class T,class Ref, class Ptr>
struct RBtreeIterator
{
	// 定义 Node 类型，它是红黑树节点的类型，这里使用了模板参数 T
	typedef RBtreeNode<T> Node;
	// 定义 self 类型，它就是当前的 RBtreeIterator 类型，方便后续使用
	typedef RBtreeIterator<T,Ref, Ptr> self;
	// 迭代器内部维护一个指向红黑树节点的指针
	Node* _node;
	Node* _root;
	//构造，初始化迭代器
	RBtreeIterator(Node* node,Node* root)
		:_node(node)
		,_root(root)
	{ }

	// 重载解引用运算符 *
	Ref operator* () //set用
	{
		return _node->_data;
	}
	// 重载箭头运算符 ->
	Ptr operator-> () //map用
	{
		//返回_data的地址
		return &_node->_data;
	}

	self operator++()
	{
		//右节点存在
		if (_node->_right)
		{
			//右节点给min
			Node* min = _node->_right;
			//循环找最左节点
			while (min->_left)
			{
				min = min->_left;
			}
			_node = min;
		}
		else //右节点不存在
		{
			//保存当前节点
			Node* cur = _node;
			//保存当前节点的父亲
			Node* parent = _node->_parent;
			//父亲节点存在，并且父亲右节点等于cur，继续往上找
			while (parent && parent->_right == cur)
			{
				cur = parent;
				parent = parent->_parent;
			}
			_node = parent;
		}
		// 返回当前迭代器对象
		return *this;
	}
	// 重载前置自减运算符 --，用于将迭代器移动到上一个节点
	self operator--()
	{
		// 若当前迭代器指向空节点
		if (_node == nullptr)
		{
			// 从根节点开始找最右节点
			Node* cur = _root;
			while (cur->_right)
			{
				cur = cur->_right;
			}
			// 迭代器指向最右节点
			_node = cur;
		}
		// 若当前节点有左子节点
		else if (_node->_left)
		{
			// 找到左子树的最右节点
			Node* min = _node->_left;
			while (min->_right)
			{
				min = min->_right;
			}
			// 迭代器指向该节点
			_node = min;
		}
		// 若当前节点无左子节点
		else
		{
			Node* cur = _node;
			Node* parent = cur->_parent;
			// 向上找第一个不是其父节点左子节点的节点
			while (parent && parent->_left == cur)
			{
				cur = parent;
				parent = parent->_parent;
			}
			// 迭代器指向该父节点
			_node = parent;
		}
		// 返回当前迭代器对象
		return *this;
	}

	// 重载不等于运算符 !=
	// 用于比较两个迭代器是否不相等
	// 比较的是两个迭代器所指向的节点是否不同
	bool operator!=(const self& k)
	{
		return _node != k._node;
	}

	// 重载等于运算符 ==
	// 用于比较两个迭代器是否相等
	// 比较的是两个迭代器所指向的节点是否相同
	bool operator==(const self& k)
	{
		return _node == k._node;
	}
};

template<class K,class T,class KeyOFT>
class RBtree
{
	typedef RBtreeNode<T> Node;
public:
	// 定义迭代器类型，使用 RBtreeIterator 模板，传入存储的数据类型 T
	typedef RBtreeIterator<T,T&,T*> Iterator;
	typedef RBtreeIterator<T,const T&,const T*> constIterator;

	Iterator begin()
	{
		Node* cur = _root;
		while (cur->_left)
		{
			cur = cur->_left;
		}
		//传cur去初始化迭代器,然后返回
		return Iterator(cur,_root);
	}

	Iterator end()
	{
		//传nullptr去初始化迭代器，然后返回
		return Iterator(nullptr,_root);
	}

	constIterator Begin() const
	{
		Node* cur = _root;
		while (cur->_left)
		{
			cur = cur->_left;
		}
		//传cur去初始化迭代器,然后返回
		return constIterator(cur, _root);
	}

	constIterator End() const 
	{
		//传nullptr去初始化迭代器，然后返回
		return constIterator(nullptr, _root);
	}

	//插入数据
	pair<Iterator,bool> insert(const T& data)
	{
		if (_root == nullptr)
		{
			_root = new Node(data);
			_root->_col = BLACK;
			return make_pair(Iterator(_root,_root),true);
		}
		Node* cur = _root;
		Node* parent = nullptr;
		KeyOFT kot; //模仿函数比较
		while (cur)
		{
			if (kot(data) < kot(cur->_data))
			{
				parent = cur;
				cur = cur->_left;
			}
			else if(kot(data) > kot(cur->_data))
			{
				parent = cur;
				cur = cur->_right;
			}
			else
			{
				// 返回一个迭代器，指向当前节点，以及一个布尔值 false 表示未插入新节点
				return make_pair(Iterator(cur, _root), false);
			}
		}
		//插入红色节点
		cur = new Node(data);
		Node* newNode = cur;
		cur->_col = RED;
		if (kot(data) < kot(parent->_data))
		{
			parent->_left = cur;
		}
		else
		{
			parent->_right = cur;
		}
		cur->_parent = parent;
		//变色+旋转
		while (parent && parent->_col == RED)
		{
			Node* grand = parent->_parent;
			if (grand->_left == parent)
			{
				Node* uncle = grand->_right;
				if (uncle && uncle->_col == RED)
				{
					parent->_col = uncle->_col = BLACK;
					grand->_col = RED;

					//继续向上更新
					cur = grand;
					parent = cur->_parent;
				}
				else
				{
					if (parent->_left == cur)
					{
						RotR(grand);
						grand->_col = RED;
						parent->_col = BLACK;
					}
					else
					{
						RotL(parent);
						RotR(grand);
						grand->_col = RED;
						cur->_col = BLACK;
					}
					//break;
				}
				
			}
			else
			{
				Node* uncle = grand->_left;
				if (uncle && uncle->_col == RED)
				{
					uncle->_col = parent->_col = BLACK;
					grand->_col = RED;

					cur = grand;
					parent = cur->_parent;
				}
				else
				{
					if (parent->_right == cur)
					{
						RotL(grand);
						grand->_col = RED;
						parent->_col = BLACK;
					}
					else
					{
						RotR(parent);
						RotL(grand);
						grand->_col = RED;
						cur->_col = BLACK;
					}
					//break;
				}
				
			}
		}
		_root->_col = BLACK;
		return make_pair(Iterator(newNode, _root), true);

	}
	//右旋
	void RotR(Node* root)
	{
		Node* subL = root->_left;
		Node* subR = subL->_right;

		root->_left = subR;
		Node* root_parent = root->_parent;
		if (subR != nullptr)
		{
			subR->_parent = root;
		}

		subL->_right = root;
		root->_parent = subL;

		if (root_parent == nullptr)
		{
			_root = subL;
			subL->_parent = nullptr;
		}
		else
		{
			if (root_parent->_left == root)
			{
				root_parent->_left = subL;
			}
			else
			{
				root_parent->_right = subL;
			}
			subL->_parent = root_parent;
		}

	}
	//左旋
	void RotL(Node* root)
	{
		Node* subL = root->_right;
		Node* subR = subL->_left;

		root->_right = subR;
		if (subR != nullptr)
		{
			subR->_parent = root;
		}
		Node* root_parent = root->_parent;

		subL->_left = root;
		root->_parent = subL;
		if (root_parent == nullptr)
		{
			_root = subL;
			subL->_parent = nullptr;
		}
		else
		{
			if (root_parent->_left == root)
			{
				root_parent->_left = subL;

			}
			else
			{
				root_parent->_right = subL;
			}
		}
	}
	//查询
	Iterator Find(const K& key)
	{
		Node* cur = _root;
		while (cur)
		{
			if (key < cur->_data.first)
			{
				cur = cur->_left;
			}
			else if(key > cur->_data.first)
			{
				cur = cur->_right;
			}
			else
			{
				return Iterator(cur, _root);
			}
		}
		return end();
	}

	void print()
	{
		_InOrder(_root);
	}
protected:
	void _InOrder(Node* root)
	{
		if (root == nullptr)
		{
			return;
		}
		_InOrder(root->_left);
		cout << root->_kv.first << ":" << root->_kv.second << endl;
		_InOrder(root->_right);
	}
private:
	Node* _root = nullptr;
};